package basic;

import java.util.Arrays;
import java.util.List;

public class RBTree {

    private TreeNode rootNode;

    public void insert(int value) {
        TreeNode treeNode = new TreeNode(value);
        if (rootNode == null) {
            treeNode.red = false;
            rootNode = treeNode;
            return;
        }
        TreeNode parentNode = rootNode;
        while (parentNode != null) {
            if (parentNode.value < value && parentNode.right != null) {
                parentNode = parentNode.right;
            } else if (parentNode.value > value && parentNode.left != null) {
                parentNode = parentNode.left;
            } else {
                break;
            }
        }

        if (parentNode.value > value) {
            parentNode.left = treeNode;
        } else if (parentNode.value < value) {
            parentNode.right = treeNode;
        }
        treeNode.parent = parentNode;
        //修复节点
        insertFixUp(treeNode);

    }

    public void insertFixUp(TreeNode node) {
        TreeNode parent;
        TreeNode gparent;
        while (((parent = node.parent) != null) && isRed(parent)) {
            gparent = parent.parent;

            if (gparent.left != null && parent == gparent.left) {
                TreeNode uncle = gparent.right;

                if (uncle != null && isRed(uncle)) {
                    setRed(gparent);
                    setBlack(parent);
                    setBlack(uncle);
                    node = gparent;
                    continue;
                }

                if (node == parent.right) {
                    rotateLR(node);
                    parent = node;
                }

                rotateLR(parent);
                setBlack(parent);
                setRed(gparent);

            } else {
                TreeNode uncle = gparent.left;

                if (uncle != null && isRed(uncle)) {
                    setRed(gparent);
                    setBlack(parent);
                    setBlack(uncle);
                    node = gparent;
                    continue;
                }

                if (node == parent.left) {
                    rotateLR(node);
                    parent = node;
                }

                rotateLR(parent);
                setBlack(parent);
                setRed(gparent);
            }
        }
        // 将根节点设为黑色
        setBlack(this.rootNode);
    }

    private void rotateLR(TreeNode treeNode) {
        TreeNode parent = treeNode.parent;
        TreeNode pparent = parent.parent;
        TreeNode right = treeNode.right;

        treeNode.right = parent;
        parent.left = right;
        if (pparent != null && parent == pparent.left) {
            pparent.left = treeNode;
        }
        if (pparent != null && parent == pparent.right) {
            pparent.right = treeNode;
        }
        if (pparent == null) {
            this.rootNode = treeNode;
        }
        treeNode.parent = pparent;
        parent.parent = treeNode;
    }

    private boolean isRed(TreeNode treeNode){
        return treeNode.red;
    }

    private boolean isBlack(TreeNode treeNode){
        return !treeNode.red;
    }

    private void setRed(TreeNode treeNode){
        treeNode.red = true;
    }

    private void setBlack(TreeNode treeNode){
        treeNode.red = false;
    }

    private void inOrder(TreeNode node) {
        if (node == null) {
            return;
        }
        inOrder(node.left);
        System.out.print(node.value + " ");
        inOrder(node.right);
    }

    private static final List<Integer> data = Arrays.asList(10, 40, 30, 60, 90, 70, 20, 50, 80);

    public static void main(String[] args) {
        RBTree RBTree = new RBTree();
        System.out.printf("== 原始数据: ");
        data.stream().forEach((item) -> {
            System.out.print(item + " ");
            RBTree.insert(item);
        });
        RBTree.inOrder(RBTree.rootNode);
    }
}
